import pandas as pd
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker
from typing import List, Dict, Any

class SQLDatabaseTool:
    """MySQL数据库操作类"""

    def __init__(self, db_config: dict[str,any], close_session: bool = True,shcema_name_str: str = None):
        """
        初始化数据库连接
        Args:
            db_url (str): 数据库连接URL，格式：mysql+pymysql://user:password@host:port/database
        """
        self.db_config = db_config

        self.engine = create_engine(create_database_url(db_config))
        self.inspector = inspect(self.engine)

        self.Session = sessionmaker(bind=self.engine)
        self.close_session  = close_session

    def is_sqlserver(self):
        return self.db_config['type'] == 'SQLSERVER'

    def get_full_table_name(self,table_name):
        if self.db_config['schema_name'] is not None:
            return self.db_config['schema_name'] + '.'+ table_name
        else:
            return table_name

    def get_table_names(self) -> List[str]:
        """获取数据库中所有表的名称"""
        self.inspector = inspect(self.engine)
        return self.inspector.get_table_names(schema=self.db_config.get('schema_name',None))


    def has_table(self, table_name: str) -> bool:
        self.inspector = inspect(self.engine)
        return self.inspector.has_table(table_name, schema=self.db_config.get('schema_name',None))

    def get_table_data(self, table_name: str, where_clause : str = None ,limit: int = 5,offset: int = 0) -> Dict[str, Any]:
        """
        获取指定表的前几行数据

        Args:
            table_name (str): 表名
            limit (int): 返回的行数，默认5行

        Returns:
            Dict[str, Any]: 包含表结构和数据的字典
        """
        session = self.Session()
        try:
            # 获取表结构
            columns = self.inspector.get_columns(table_name,schema=self.db_config.get('schema_name',None))
            column_names = [col['name'] for col in columns]
            full_table_name = table_name
            if self.db_config.get('schema_name',None) is not None:
                full_table_name = self.db_config['schema_name'] + '.' + table_name
            if self.db_config['type'] == 'SQLSERVER':
                if limit > 0:
                    base_query = f"SELECT TOP {limit} * FROM {full_table_name}"
                else:
                    base_query = f"SELECT * FROM {full_table_name}"
            else:
                base_query = f"SELECT * FROM {full_table_name}"

            if where_clause:  # 当 WHERE 条件存在时
                base_query += f" WHERE {where_clause}"

            if limit > 0 and self.db_config['type'] != 'SQLSERVER':
                base_query += f" LIMIT {limit}  OFFSET {offset} "
            # 查询数据
            result = session.execute(text(base_query)).fetchall()

            # 将结果转换为字典列表
            data = []
            for row in result:
                row_dict = {}
                for i, value in enumerate(row):
                    row_dict[column_names[i]] = value  # 将所有值转换为字符串，确保JSON序列化
                data.append(row_dict)

            return {
                "table_name": table_name,
                "columns": column_names,
                "data": data,
                "total_rows": len(data)
            }

        except Exception as e:
            return {
                "error": f"Error fetching data: {str(e)}"
            }
        finally:
            if self.close_session:
                session.close()

    def insert_data(self, table_name: str, data: Dict[str, Any]) -> bool:
        """
        插入数据到指定表

        Args:
            table_name (str): 表名
            data (Dict[str, Any]): 要插入的数据字典，键为列名，值为对应值

        Returns:
            bool: 是否成功插入
        """
        session = self.Session()
        try:
            # 构建插入语句
            columns = ', '.join(data.keys())
            placeholders = ', '.join([':{}'.format(key) for key in data.keys()])
            print(columns)
            # 执行插入
            session.execute(text(f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"), data)
            session.commit()
            return True

        except Exception as e:
            session.rollback()
            print(e)
            return False

        finally:
            if self.close_session:
                session.close()

    def update_data(self, table_name: str, data: Dict[str, Any], 
                    where_clause: str,params: Dict[str, Any]) -> bool:
        """
        更新指定表中的数据

        Args:
            table_name (str): 表名
            data (Dict[str, Any]): 要更新的数据字典，键为列名，值为对应值
            where_clause (str): WHERE子句
            params (Dict[str, Any]): 参数字典，包含更新数据和WHERE子句参数

        Returns:
            bool: 是否成功更新
        """
        session = self.Session()
        try:
            # 构建更新语句
            set_clause = ', '.join([f"{col} = :{col}" for col in data.keys()])

            print(text(f"UPDATE {table_name} SET {set_clause} WHERE {where_clause}"))
            # 执行更新
            session.execute(
                text(f"UPDATE {table_name} SET {set_clause} WHERE {where_clause}"),params
            )
            session.commit()
            return True

        except Exception as e:
            session.rollback()
            print(e)
            return False

        finally:
            if self.close_session:
                session.close()

    def delete_data(self,table_name: str, where_clause: str ) -> bool:
        """
        从指定表中删除数据

        Args:
            table_name (str): 表名
            where_clause (str): WHERE子句
            params (Dict[str, Any]): 参数字典

        Returns:
            bool: 是否成功删除
        """
        session = self.Session()
        try:
            # 执行删除
            session.execute(
                text(f"DELETE FROM {table_name} WHERE {where_clause}")
            )
            session.commit()
            return True

        except Exception as e:
            session.rollback()
            return False

        finally:
            if self.close_session:
                session.close()


    def count(self, table_name: str) -> int:
        session = self.Session()
        try:

            # 执行更新
            result = session.execute(
                text(f"select count(*) from {self.db_config['schema_name']+'.' if self.db_config['schema_name'] is not None else ''}{table_name}")
            )

            return result.scalar()

        except Exception as e:
            print(e)
            return -1

        finally:
            if self.close_session:
                session.close()

    def get_table_data_frame(self, table_name, size=-1,offset=0):
        full_table_name = table_name
        if self.db_config['schema_name'] is not None:
            full_table_name = self.db_config['schema_name'] + '.'+ table_name

        if size < 0:
            sql = f"select * from {full_table_name}"
        else:
            if self.db_config['type'] == 'SQLSERVER':
                sql = f"select top {size} * from {full_table_name}"
            else:
                sql = f"select * from {self.db_config['schema_name']}.{table_name} LIMIT {size} OFFSET {offset}"

        return  pd.read_sql_query(sql, self.engine)


    def truncate_table(self, table_name: str):
        session = self.Session()
        try:
            if self.db_config['schema_name'] is not None:
                table_name = self.db_config['schema_name'] + '.' + table_name
            if self.db_config['type'] == 'PostgreSQL':
                session.execute(text(f"TRUNCATE TABLE {table_name} RESTART IDENTITY"))
            else:
                session.execute(text(f"TRUNCATE TABLE {table_name} "))
            session.commit()
            return True

        except Exception as e:
            session.rollback()
            print(e)
            return False

        finally:
            if self.close_session:
                session.close()

    def execute_sql(self, sql: str):
        session = self.Session()
        try:
            session.execute(text(sql))
            session.commit()
            return True

        except Exception as e:
            session.rollback()
            print(e)
            return False

        finally:
            if self.close_session:
                session.close()

def create_database_url(config):
    if config["type"] == "MySQL":
        return (f"mysql+pymysql://{config['username']}:{config['password']}@"
                f"{config['host']}:{config['port']}/{config['database']}")
    elif config["type"] == "PostgreSQL":
        return (f"postgresql://{config['username']}:{config['password']}@"
                f"{config['host']}:{config['port']}/{config['database']}")
    elif config['type'] == 'Sqlite':
        return config.get('url')
    elif config['type'] == 'SQLSERVER':
        return (f"mssql+pymssql://{config['username']}:{config['password']}@"
                f"{config['host']}:{config['port']}/{config['database']}")
    else:
        return (f"postgresql://{config['username']}:{config['password']}@"
                f"{config['host']}:{config['port']}/{config['database']}")





